package correlation

import (
	"fmt"
	"math"

	"gonum.org/v1/gonum/floats"
	"gonum.org/v1/gonum/mat"
	"gonum.org/v1/plot"
	"gonum.org/v1/plot/plotter"
)

func CorrMatrix(matrix *mat.Dense) *mat.SymDense {

	matrixT := mat.DenseCopyOf(matrix.T())
	ret := mat.NewSymDense(matrix.RawMatrix().Cols, nil)

	//подсчет корреляции
	for i := 0; i < matrix.RawMatrix().Cols; i++ {
		x := matrixT.RawRowView(i)
		xS := floats.Sum(x) / float64(len(x))
		for j := i; j < matrix.RawMatrix().Cols; j++ {
			y := matrixT.RawRowView(j)
			yS := floats.Sum(y) / float64(len(y))

			var Sxy, Sx, Sy float64

			// корреляция столбцов i и j
			for k := 0; k < len(x); k++ {
				tmpX, tmpY := x[k]-xS, y[k]-yS

				Sxy += tmpX * tmpY
				Sx += math.Pow(tmpX, 2)
				Sy += math.Pow(tmpY, 2)
			}
			kor := Sxy / (math.Sqrt(Sx) * math.Sqrt(Sy))
			ret.SetSym(i, j, kor)
		}
	}
	/* 	corrM := mat.SymDense{}
	   	stat.CorrelationMatrix(&corrM, matrix, nil)
	*/
	/* 	fmt.Println("\n------\n", ret.RawSymmetric().Data)
	   	fmt.Println("\n------\n", corrM.RawSymmetric().Data)
	   	fmt.Println("\n------\n") */

	return ret
}

func PrintCorrMatrix(matrix *mat.SymDense, names []string) {
	fmt.Println("Correlation Matrix:")
	N := matrix.RawSymmetric().N

	sizes := make([]int, 0)
	hNames := ""
	for _, name := range names {
		hNames += fmt.Sprint(" ", name)
		sizes = append(sizes, len(name)+1)
	}
	fmt.Println(hNames)

	for i := 0; i < N; i++ {
		for j := 0; j < N; j++ {
			fmt.Printf(fmt.Sprintf("%% %d.3f", sizes[j]), matrix.At(i, j))
			if j == N-1 {
				fmt.Print(" ", names[i])
			}
		}
		fmt.Println()
	}
	fmt.Printf("\n")
}

func CreateHeatMap(matrix *mat.SymDense) {
	//Создание графика
	p := plot.New()
	heatmap := plotter.NewHeatMap(NewMyGridXYZ(*matrix), myPalette{})
	p.Add(heatmap)
	p.HideX()
	p.HideY()

	//h := plotter.NewHeatMap(plotter.NewHeatMapGrid, palette.Heat(256, 1))

	if err := p.Save(500, 500, "out/heatmap.png"); err != nil {
		fmt.Println("Ошибка сохранения тепловой карты:", err)
		return
	}
}
